import time, sys, os, argparse
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
import torchvision.transforms as transforms

from SkyNet import *

def convert2cpu(gpu_matrix):
    return torch.FloatTensor(gpu_matrix.size()).copy_(gpu_matrix)

class DACDataset(Dataset):
    def __init__(self, root, shape=None, transform=None, batch_size=32, num_workers=4):
        self.files = [file for file in os.listdir(root) if os.path.isfile(os.path.join(root, file))]
        self.imageNames = [file.split('.')[0] for file in self.files]
        self.files = [os.path.join(root, file) for file in self.files]
        self.nSamples = len(self.files)
        self.transform = transform
        self.shape = shape
        self.batch_size = batch_size
        self.num_workers = num_workers
       
    def __len__(self):
        return self.nSamples

    def __getitem__(self, index):
        imgpath = self.files[index]
        img = Image.open(imgpath).convert('RGB')
        if self.shape:
            img = img.resize(self.shape)

        if self.transform is not None:
            img = self.transform(img)

        return img

def qdata(x, SA, na):
    nB = x.size(0)
    nC = x.size(1)
    nH = x.size(2)
    nW = x.size(3)

    amin = -2.**(na - 1)
    amax = 2.**(na - 1) - 1

    QSA = SA.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand(nB,nC,nH,nW).cuda()
    qx = x.div(QSA).clamp(amin, amax).round().cuda()
    return qx

def qconv(x, w, b ,m, na, nb, nm, relu = True, last = False):
    nB = x.size(0)
    nI = x.size(1)
    nH = x.size(2)
    nW = x.size(3)
    nO = w.size(0)
    nG = int(x.size(1)/w.size(1))
    nK = w.size(2)
    
    conv2d = nn.Conv2d(nI, nO, nK, 1, int((nK-1)/2), groups=nG, bias=False)
    
    conv2d.weight = nn.Parameter(w)
    qb = b.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand(nB,nO,nH,nW)
    qm = m.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand(nB,nO,nH,nW)
    
    qy = conv2d(x)
    qy = qy.clamp(-2.**(nb-1), 2.**(nb-1)-1)
    qy = qy + qb
    if(relu):
        qy = nn.ReLU()(qy)
        y = qy.mul(qm).div(2**nm)
        if(not last):
            y = y.clamp(0, 2.**na-1).round()
            y = y.round()
    else:
        y = qy.mul(qm).div(2**nm)
        if(not last):
            y = y.clamp(-2.**(na-1), 2.**(na-1) - 1).round()
    return y

def qinference(x, qconv_layers, qbn_layers, na, nb, nm):
    x = qdata(x, torch.Tensor([0.015625,0.015625,0.015625]), na)
    x = qconv(x, qconv_layers[0].weight, qconv_layers[0].bias, qbn_layers[0].weight, na, nb, nm)
    x = qconv(x, qconv_layers[1].weight, qconv_layers[1].bias, qbn_layers[1].weight, na, nb, nm)
    x = nn.MaxPool2d(kernel_size=2, stride=2)(x)
    x = qconv(x, qconv_layers[2].weight, qconv_layers[2].bias, qbn_layers[2].weight, na, nb, nm)
    x = qconv(x, qconv_layers[3].weight, qconv_layers[3].bias, qbn_layers[3].weight, na, nb, nm)
    x = nn.MaxPool2d(kernel_size=2, stride=2)(x)
    x = qconv(x, qconv_layers[4].weight, qconv_layers[4].bias, qbn_layers[4].weight, na, nb, nm)
    x = qconv(x, qconv_layers[5].weight, qconv_layers[5].bias, qbn_layers[5].weight, na, nb, nm)

    reorg = ReorgLayer()(x)
    
    x = nn.MaxPool2d(kernel_size=2, stride=2)(x)
    x = qconv(x, qconv_layers[6].weight, qconv_layers[6].bias, qbn_layers[6].weight, na, nb, nm)
    x = qconv(x, qconv_layers[7].weight, qconv_layers[7].bias, qbn_layers[7].weight, na, nb, nm)
    x = qconv(x, qconv_layers[8].weight, qconv_layers[8].bias, qbn_layers[8].weight, na, nb, nm)
    x = qconv(x, qconv_layers[9].weight, qconv_layers[9].bias, qbn_layers[9].weight, na, nb, nm)

    x = torch.cat([reorg,x],1)

    x = qconv(x, qconv_layers[10].weight, qconv_layers[10].bias, qbn_layers[10].weight, na, nb, nm)
    x = qconv(x, qconv_layers[11].weight, qconv_layers[11].bias, qbn_layers[11].weight, na, nb, nm)
    x = qconv(x, qconv_layers[12].weight, qconv_layers[12].bias, qbn_layers[12].weight, na, nb, nm, relu = False, last = True)

    return x

imgDir = '/media/Workspace/DAC/val/jpg'
DEVICE = torch.device("cuda:0")

parser = argparse.ArgumentParser(description="SkyNet Data Free Quantization on DAC dataset.")
parser.add_argument("--nw", type=int, default=6)
parser.add_argument("--na", type=int, default=8)
parser.add_argument("--nb", type=int, default=16)
parser.add_argument("--nm", type=int, default=16)
parser.add_argument("--batch_size", type=int, default=16)
parser.add_argument('--workers', type=int, default=16, metavar='N')
args = parser.parse_args()
print(args)

if __name__ == '__main__':
    net = SkyNet()
    net.load_state_dict(torch.load('./QSkyNet.pth'))
    net = net.to(DEVICE)
    dataset = DACDataset(imgDir, shape=(net.width, net.height),
                        transform=transforms.Compose([
                                transforms.ToTensor(),
                                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.25, 0.25, 0.25]),]))
    test_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.workers, pin_memory=True)

    net = net.cuda()
    net.eval()
    anchors = net.anchors
    num_anchors = net.num_anchors
    anchor_step = len(anchors) // num_anchors
    h = 20
    w = 40
    total = 0
    imageNum = dataset.__len__()
    results = np.zeros((imageNum, 5))

    conv_layers = []
    bn_layers = []
    for module in net.named_modules():
        if isinstance(module[1], nn.Conv2d):
            conv_layers.append(module[1])
        if isinstance(module[1], nn.BatchNorm2d):
            bn_layers.append(module[1])

    grid_x = torch.linspace(0, w - 1, w).repeat(h, 1).repeat(args.batch_size * num_anchors, 1, 1).view(args.batch_size * num_anchors * h * w).cuda()
    grid_y = torch.linspace(0, h - 1, h).repeat(w, 1).t().repeat(args.batch_size * num_anchors, 1, 1).view(args.batch_size * num_anchors * h * w).cuda()
    anchor_w = torch.Tensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([0]))
    anchor_h = torch.Tensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([1]))
    anchor_w = anchor_w.repeat(args.batch_size, 1).repeat(1, 1, h * w).view(args.batch_size * num_anchors * h * w).cuda()
    anchor_h = anchor_h.repeat(args.batch_size, 1).repeat(1, 1, h * w).view(args.batch_size * num_anchors * h * w).cuda()
    sz_hw = h * w
    sz_hwa = sz_hw * num_anchors

    stime = time.time()
    for batch_idx, data in enumerate(test_loader):
        data = data.cuda()
        output = qinference(data, conv_layers, bn_layers, args.na, args.nb, args.nm)
        batch = output.size(0)
        output = output.view(batch * num_anchors, 5, h * w).transpose(0, 1).contiguous().view(5, batch * num_anchors * h * w)

        if batch != args.batch_size:
            grid_x = torch.linspace(0, w - 1, w).repeat(h, 1).repeat(batch * num_anchors, 1, 1).view(batch * num_anchors * h * w).cuda()
            grid_y = torch.linspace(0, h - 1, h).repeat(w, 1).t().repeat(batch * num_anchors, 1, 1).view(batch * num_anchors * h * w).cuda()
            anchor_w = torch.Tensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([0]))
            anchor_h = torch.Tensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([1]))
            anchor_w = anchor_w.repeat(batch, 1).repeat(1, 1, h * w).view(batch * num_anchors * h * w).cuda()
            anchor_h = anchor_h.repeat(batch, 1).repeat(1, 1, h * w).view(batch * num_anchors * h * w).cuda()

        det_confs = torch.sigmoid(output[4])
        det_confs = convert2cpu(det_confs)

        for b in range(batch):
            det_confs_inb = det_confs[b * sz_hwa:(b + 1) * sz_hwa].detach().numpy()
            ind = np.argmax(det_confs_inb)

            xs_inb = torch.sigmoid(output[0, b * sz_hwa + ind]) + grid_x[b * sz_hwa + ind]
            ys_inb = torch.sigmoid(output[1, b * sz_hwa + ind]) + grid_y[b * sz_hwa + ind]
            ws_inb = torch.exp(output[2, b * sz_hwa + ind]) * anchor_w[b * sz_hwa + ind]
            hs_inb = torch.exp(output[3, b * sz_hwa + ind]) * anchor_h[b * sz_hwa + ind]

            bcx = xs_inb.item() / w
            bcy = ys_inb.item() / h
            bw = ws_inb.item() / w
            bh = hs_inb.item() / h

            xmin = bcx - bw / 2.0
            ymin = bcy - bh / 2.0
            xmax = xmin + bw
            ymax = ymin + bh

            results[total + b, 1:] = np.asarray([xmin * 640, xmax * 640, ymin * 360, ymax * 360])

        total += batch
    etime = time.time()

    results[:, 0] = dataset.imageNames
    index = np.argsort(results[:, 0])

    file = open('./result.txt','w+')
    for i in range(len(index)):
        name = '%03d'%int(results[index[i]][0])+'.jpg '
        bbox = '['+str(int(round(results[index[i]][1])))+', '+str(int(round(results[index[i]][2])))+', '+str(int(round(results[index[i]][3])))+', '+str(int(round(results[index[i]][4])))+']\n'
        file.write(name+bbox)
    file.close()